from solvers.solverbase import SolverBase
from solvers.supg.supgbase import SUPGBase
from solvers.stabilizations import supg
from dolfin import *
import numpy

class GaussLegendre4(SolverBase,SUPGBase):
    '2-symplectic method rewritten without increments. SUPG stabilized.'
    def __init__(self,p=1):
        SolverBase.__init__(self,p)
        SUPGBase.__init__(self)

    def solve(self,problem):
        mesh = problem.domain

        V = FunctionSpace(mesh,'CG',1)
        VV = VectorFunctionSpace(mesh,'CG',1)
        DG = FunctionSpace(mesh,'DG',0)

        phi_ = interpolate(problem.phi_,V)
        phiP = Function(V)                  # solution to the auxiliary equation
        phiC = Function(V)                  # solution to the true equation
        phi = TrialFunction(V)          
        varphi = TestFunction(V)

        t = problem.t          
        dt = Constant(SolverBase._get_time_step(self,problem,VV,t))
        T = problem.T

        v = SolverBase._get_velocity(self,problem,VV,t)
        tau = Function(DG)
        
        tau.vector()[:] = supg.get_tau(v,eps=supg.diffusivity)  # stab. parameters

        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                           ) # a way to realize swith on |v| == 0.
 
        aP = phi*varphi*dx #+ tau*pert*phi*inner(v,nabla_grad(varphi))*dx

        lP = phi_*varphi*dx\
             + dt*sqrt(3.)/6.*inner(v,nabla_grad(phiC))*varphi*dx \
             - dt*sqrt(3.)/6.*inner(v,nabla_grad(phi_))*varphi*dx #\
            #+ tau*pert*phi_*inner(v,nabla_grad(varphi))*dx \


            # + dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiC),outer(v,v)),nabla_grad(varphi))*dx \
            # - dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx

        aC = phi*varphi*dx + 0.5*dt*inner(v,nabla_grad(phi))*varphi*dx\
             +tau*pert*phi*inner(v,nabla_grad(varphi))*dx + 0.5*dt*tau*pert*inner(dot(nabla_grad(phi),outer(v,v)),nabla_grad(varphi))*dx 

        lC = phi_*varphi*dx + tau*pert*phi_*inner(v,nabla_grad(varphi))*dx\
             -dt*(0.5-sqrt(3.)/6.)*inner(v,nabla_grad(phi_))*varphi*dx\
             -dt*(0.5-sqrt(3.)/6.)*tau*pert*inner(dot(nabla_grad(phi_),outer(v,v)),nabla_grad(varphi))*dx\
             -dt*sqrt(3.)/6.*inner(v,nabla_grad(phiP))*varphi*dx\
             -dt*sqrt(3.)/6.*tau*pert*inner(dot(nabla_grad(phiP),outer(v,v)),nabla_grad(varphi))*dx

        self.update(problem,phi_,t,float(dt),0)  # store the solution at time 0

        phiC.assign(phi_)  # initial guess
        
        phiC_ = Function(V) # so that error could be calculated
        phiC_.assign(phi_)

        AP = assemble(aP)
        AC = assemble(aC)
        phi = Function(V)
   
        t = 0.
        count = 0
       
        maxIter = 10
        tol = 1E-8
        
        while t <= T:
            t += float(dt)
            count += 1
            error = 1.
            iteration = 0

            while error > tol and iteration < maxIter:
                iteration += 1
                LP = assemble(lP)
                solve(AP,phi.vector(),LP)  # solve aux system
                phiP.assign(phi)

                LC = assemble(lC)
                solve(AC,phi.vector(),LC)  # solve true system
                phiC.assign(phi)

                diff = phiC.vector().array() - phiC_.vector().array()
                error = numpy.linalg.norm(diff,ord=2)        # compute the error
                phiC_.assign(phiC)
                
        #        print iteration, error

            phi_.assign(phiC)
            self.update(problem,phi_,t,float(dt),count)

            
        problem.save_data(self.saveDir,str(self.CFL))
